package com.hsyt.pfc.support.filter;

import com.alibaba.dubbo.common.Constants;
import com.alibaba.dubbo.common.extension.Activate;
import com.alibaba.dubbo.rpc.*;
import com.hsyt.pfc.support.util.TraceIdUtil;
import org.apache.commons.lang3.StringUtils;
import org.apache.logging.log4j.ThreadContext;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.UUID;

@Activate(group = {Constants.CONSUMER, Constants.PROVIDER})
public class TraceFilter implements Filter {

    private static final Logger log = LoggerFactory.getLogger(TraceFilter.class);
    private final String TRACE_KEY = "dubbo.traceid";
    private final String GROUP_KEY = "dubbo.groupName";

    @Override
    public Result invoke(Invoker<?> invoker, Invocation invocation)
            throws RpcException {

        RpcContext rpcContext = RpcContext.getContext();
        String traceId = rpcContext.getAttachment(TRACE_KEY);
        String groupName = rpcContext.getAttachment(GROUP_KEY);
        if (StringUtils.isBlank(traceId)) {
            traceId = TraceIdUtil.getLocalTraceId();

            if (traceId == null) {
                traceId = UUID.randomUUID().toString().replace("-", "");
            }
            rpcContext.setAttachment(TRACE_KEY, traceId);
            setLocalTraceId(traceId);
        } else {
            setLocalTraceId(traceId);
        }
        if (StringUtils.isBlank(groupName)) {
            groupName = invocation.getAttachment("group");
            if (StringUtils.isNotBlank(groupName)) {
                rpcContext.setAttachment(GROUP_KEY, groupName);
                setLocalGroupName(groupName);
            }
        } else {
            setLocalGroupName(groupName);
        }
        Result result = invoker.invoke(invocation);
        return result;


    }

    public static void setLocalTraceId(String traceId) {
        TraceIdUtil.setLocalTraceId(traceId);
        ThreadContext.put("logId", traceId);
    }

    private static void setLocalGroupName(String groupName) {
        ThreadContext.put("groupName", groupName);
    }
}

